import copy
import json
import logging
import numpy as np
import torch.utils.data as Data

logger = logging.getLogger(__name__)


def read_csv(path):
    """
    读取CSV文件内容到列表

    参数:
        path: 文件路径

    返回:
        data_list: 包含文本对和标签的列表，格式为 [[text_a, text_b, label], ...]
    """
    data_list = []
    with open(path, 'r', encoding='UTF-8') as input_data:
        # 加载JSON内容
        json_content = json.load(input_data)
        # 逐条读取记录
        for block in json_content:
            # 获取文本对和标签
            text_a = block['query1']
            text_b = block['query2']
            label = block['label']
            # 检查标签是否合法，不合法则设为0
            if label not in ['0', '1', '2']:
                label = 0
            data_list.append([text_a, text_b, int(label)])
    return data_list


def get_dateset_labels(path):
    """
    从数据文件中提取所有标签

    参数:
        path: 文件路径

    返回:
        label_list: 标签列表
    """
    label_list = []
    with open(path, "r", encoding="utf-8") as input_data:
        json_content = json.load(input_data)
        # 逐条读取记录
        for block in json_content:
            label = block['label']
            # 检查标签是否合法，不合法则设为0
            if label not in ['0', '1', '2']:
                label = 0
            label_list.append(int(label))
    return label_list


class InputFeatures(object):
    """
    单个数据的特征集合

    属性:
        input_ids: 输入token的ID列表
        attention_mask: 注意力掩码，用于指示哪些token需要被关注
        token_type_ids: token类型ID，用于区分不同的句子
        label: 样本标签
    """

    def __init__(self, input_ids, attention_mask=None, token_type_ids=None, label=None):
        self.input_ids = input_ids
        self.attention_mask = attention_mask
        self.token_type_ids = token_type_ids
        self.label = label

    def __repr__(self):
        return str(self.to_json_string())

    def to_dict(self):
        """将实例序列化为Python字典"""
        output = copy.deepcopy(self.__dict__)
        return output

    def to_json_string(self):
        """将实例序列化为JSON字符串"""
        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"


def convert_examples_to_features(
        examples,
        tokenizer,
        max_length=512,
        pad_token=0,
        pad_token_segment_id=0,
        data_type=None,
):
    """
    将文本示例转换为模型可以接受的特征格式

    参数:
        examples: 输入示例列表，格式为 [[sentences1, sentences2, label], ...]
        tokenizer: 分词器实例
        max_length: 最大序列长度
        pad_token: 填充token的ID
        pad_token_segment_id: 填充token的类型ID

    返回:
        features: 特征列表，每个元素是一个InputFeatures实例
    """
    features = []
    for example in examples:
        # 使用分词器对文本对进行编码，添加特殊token并截断到最大长度
        inputs = tokenizer.encode_plus(example[0], example[1], add_special_tokens=True, max_length=max_length)
        input_ids, token_type_ids = inputs["input_ids"], inputs["token_type_ids"]
        # 初始化注意力掩码，所有有效token位置为1
        attention_mask = [1] * len(input_ids)

        # 计算需要填充的长度
        padding_length = max_length - len(input_ids)

        # 对序列进行填充，使所有序列长度一致
        input_ids = input_ids + ([pad_token] * padding_length)
        attention_mask = attention_mask + ([0] * padding_length)  # 填充位置的注意力掩码为0
        token_type_ids = token_type_ids + ([pad_token_segment_id] * padding_length)
        label = example[2]

        # 创建特征对象并添加到特征列表
        features.append(
            InputFeatures(input_ids, attention_mask, token_type_ids, label)
        )
    return features


class BuildDataSet(Data.Dataset):
    """
    将经过特征转换的数据包装成PyTorch Dataset

    属性:
        features: 特征列表
    """

    def __init__(self, features):
        self.features = features

    def __getitem__(self, index):
        """获取单个样本"""
        feature = self.features[index]
        # 将特征转换为numpy数组
        input_ids = np.array(feature.input_ids)
        attention_mask = np.array(feature.attention_mask)
        token_type_ids = np.array(feature.token_type_ids)
        label = np.array(feature.label)

        return input_ids, attention_mask, token_type_ids, label

    def __len__(self):
        """返回数据集大小"""
        return len(self.features)


def config_to_dict(config):
    """
    将配置对象转换为字典

    参数:
        config: 配置对象

    返回:
        output: 配置字典
    """
    output = copy.deepcopy(config.__dict__)
    if hasattr(config.__class__, "model_type"):
        output["model_type"] = config.__class__.model_type
    output['device'] = config.device.type
    return output


def config_to_json_string(config):
    """
    将配置对象序列化为JSON字符串

    参数:
        config: 配置对象

    返回:
        json_string: JSON格式的配置字符串
    """
    return json.dumps(config_to_dict(config), indent=2, sort_keys=True) + '\n'